"""
evacuation_project.py

Advanced 3D super high-rise building evacuation simulation,
with dynamic 3D stair-only visualization, dynamic floor-plan overlays,
distinct floor layouts, adjustable stair vs floor speeds,
and real-time remaining-agent statistics, saving to MP4.
Based on Ge et al. (2024) "Research on Three-Dimensional Super High-Rise Building Evacuation Modeling Considering Physical Exertion"
DOI: 10.12677/mos.2024.136540
"""

import numpy as np
import heapq
import math
import random
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from matplotlib.patches import Patch
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import animation

# -----------------------------
# 1. 参数设置
# -----------------------------
CELL_SIZE         = 1.0    # 元胞边长 (m)
FLOORS            = 10     # 楼层数
HEIGHT, WIDTH     = 30, 30 # 每层网格尺寸
AGENTS_PER_FLOOR  = 40     # 每层人数
SPEED_MEAN        = 1.0    # m/s，地面初速度均值
SPEED_STD         = 0.1    # m/s，速度标准差
STAIR_SPEED_FACTOR= 0.7    # 楼梯移动速度相对地面速度系数
FLOOR_HEIGHT      = 3.0    # 每层楼高 (m)
MAX_STEPS         = 1000   # 最大模拟步数

# -----------------------------
# 2. 楼梯与出口位置
# -----------------------------
random.seed(42)
STAIR_POSITIONS = random.sample(
    [(y, x) for y in range(2, HEIGHT-2) for x in range(2, WIDTH-2)],
    k=4
)
# 地面层出口：楼梯加两个角落
EXIT_COORDS = STAIR_POSITIONS + [(0, 0), (HEIGHT-1, WIDTH-1)]
EXIT_POSITIONS = [(0, y, x) for (y, x) in EXIT_COORDS]

# -----------------------------
# 3. 生成不同楼层布局
# -----------------------------

def build_environment():
    env = np.ones((FLOORS, HEIGHT, WIDTH), dtype=int)
    for z in range(FLOORS):
        grid = np.zeros((HEIGHT, WIDTH), dtype=int)
        pattern = z % 4
        if pattern == 0:
            grid[HEIGHT//2-1:HEIGHT//2+2, :] = 0
            grid[:, WIDTH//2-1:WIDTH//2+2] = 0
        elif pattern == 1:
            for i in range(HEIGHT):
                grid[i, i] = 0
                grid[i, WIDTH-1-i] = 0
        elif pattern == 2:
            grid[2:HEIGHT-2, 2] = 0
            grid[2:HEIGHT-2, WIDTH-3] = 0
            grid[2, 2:WIDTH-2] = 0
            grid[HEIGHT-3, 2:WIDTH-2] = 0
        else:
            grid[5::10, :] = 0
            grid[:, 5::10] = 0
        np.random.seed(z)
        for _ in range(5):
            ry = np.random.randint(1, HEIGHT-4)
            rx = np.random.randint(1, WIDTH-4)
            grid[ry:ry+3, rx:rx+3] = 1
        # 0可通行,1墙，翻转到 env 编号
        floor = np.where(grid==1, 0, 1)
        for (sy, sx) in STAIR_POSITIONS:
            floor[sy, sx] = 2
        env[z] = floor
    return env

# -----------------------------
# 4. 多源 Dijkstra 静态场
# -----------------------------

def compute_static_field(env):
    f, h, w = env.shape
    static = np.full(env.shape, np.inf)
    visited = np.zeros(env.shape, bool)
    pq = []
    for pos in EXIT_POSITIONS:
        static[pos] = 0.0
        heapq.heappush(pq, (0.0, pos))
    dirs = [(dz, dy, dx)
            for dz in (-1,0,1) for dy in (-1,0,1) for dx in (-1,0,1)
            if not (dz==dy==dx==0)]
    while pq:
        dist_curr, (z, y, x) = heapq.heappop(pq)
        if visited[z, y, x]:
            continue
        visited[z, y, x] = True
        for dz, dy, dx in dirs:
            nz, ny, nx = z+dz, y+dy, x+dx
            if 0 <= nz < f and 0 <= ny < h and 0 <= nx < w and env[nz, ny, nx] != 1:
                step = CELL_SIZE * math.sqrt(1 + (FLOOR_HEIGHT/CELL_SIZE)**2) if dz != 0 else CELL_SIZE
                nd = dist_curr + step
                if nd < static[nz, ny, nx]:
                    static[nz, ny, nx] = nd
                    heapq.heappush(pq, (nd, (nz, ny, nx)))
    return static

# -----------------------------
# 5. Agent 类：不同区域速度
# -----------------------------

class Agent:
    def __init__(self, z, y, x):
        self.pos   = (z, y, x)
        self.V0    = random.gauss(SPEED_MEAN, SPEED_STD)
        self.d_sum = 0.0
        self.t_sum = 0.0
        self.count = 0.0
    def speed(self, dz):
        base = self.V0 * (1 - 0.01 * (self.d_sum**0.71) /
                          (1 - 0.01 * (self.t_sum**0.51) * math.exp(0.01*self.t_sum)))
        return base * (STAIR_SPEED_FACTOR if dz != 0 else 1.0)

# -----------------------------
# 6. 初始化 Agents
# -----------------------------

def init_agents(env):
    agents = []
    for z in range(FLOORS):
        empties = list(zip(*np.where(env[z] == 0)))
        random.shuffle(empties)
        for y, x in empties[:AGENTS_PER_FLOOR]:
            agents.append(Agent(z, y, x))
    return agents

# -----------------------------
# 7. 模拟循环：返回帧和统计数据
# -----------------------------

def simulate(env, static):
    agents = init_agents(env)
    frames = []
    stats  = []  # (total, [per_floor])
    dirs = [(dz, dy, dx)
            for dz in (-1,0,1) for dy in (-1,0,1) for dx in (-1,0,1)
            if not (dz==dy==dx==0)]
    for step in range(MAX_STEPS):
        occ = set(a.pos for a in agents)
        positions = []
        for a in agents[:]:
            a.t_sum += 1
            a.count += 1
            z, y, x = a.pos
            tau = CELL_SIZE / a.speed(0)
            if a.count >= tau:
                cand = []
                for dz, dy, dx in dirs:
                    nz, ny, nx = z+dz, y+dy, x+dx
                    if (0 <= nz < FLOORS and 0 <= ny < HEIGHT and 0 <= nx < WIDTH
                        and env[nz, ny, nx] != 1 and (nz, ny, nx) not in occ
                        and static[nz, ny, nx] <= static[z, y, x]):
                        cand.append((static[nz, ny, nx], (dz, dy, dx)))
                if cand:
                    cand.sort(key=lambda c: c[0])
                    best = cand[0][0]
                    moves = [mv for d, mv in cand if d == best]
                    dz, dy, dx = random.choice(moves)
                    new = (z+dz, y+dy, x+dx)
                    occ.remove(a.pos)
                    occ.add(new)
                    a.pos = new
                    dist = CELL_SIZE * math.sqrt(1 + (FLOOR_HEIGHT/CELL_SIZE)**2) if dz != 0 else CELL_SIZE
                    a.d_sum += dist
                    a.count = 0
                    if new in EXIT_POSITIONS:
                        occ.remove(new)
                        agents.remove(a)
                else:
                    a.count = 0
            positions.append(a.pos)
        total = len(agents)
        per = [sum(1 for p in positions if p[0] == f) for f in range(FLOORS)]
        stats.append((total, per))
        frames.append(positions)
        if not agents:
            break
    return frames, stats

# -----------------------------
# 8. 可视化 & 保存 MP4
# -----------------------------

def visualize(env, frames, stats):
    cmap = ListedColormap(['black', 'white', 'saddlebrown'])
    fig = plt.figure(figsize=(18, 9))
    gs = fig.add_gridspec(2, 6)
    # 3D 楼梯视图
    ax3d = fig.add_subplot(gs[:, 0], projection='3d')
    ax3d.set_xlim(0, WIDTH)
    ax3d.set_ylim(0, HEIGHT)
    ax3d.set_zlim(0, FLOORS)
    vox_stairs = (env == 2)
    ax3d.voxels(vox_stairs, facecolors='saddlebrown', edgecolor='k', alpha=0.7)
    scat3d = ax3d.scatter([], [], [], s=20, c='blue')
    text_total = ax3d.text2D(0.02, 0.95, '', transform=ax3d.transAxes)
    text_floor = ax3d.text2D(0.02, 0.85, '', transform=ax3d.transAxes)
    ax3d.set_title('3D Stair Movement')

    # 平面子图
    axs = []
    scat2d = []
    for f in range(FLOORS):
        row = 0 if f < 5 else 1
        col = (f % 5) + 1
        ax = fig.add_subplot(gs[row, col])
        ax.imshow(env[f], cmap=cmap, origin='lower')
        ax.set_title(f'Floor {f}')
        ax.axis('off')
        axs.append(ax)
        scat2d.append(ax.scatter([], [], s=10, c='red'))

    def update(i):
        pts = frames[i]
        # 更新3D
        if pts:
            zs, ys, xs = zip(*pts)
            scat3d._offsets3d = (xs, ys, zs)
        else:
            scat3d._offsets3d = ([], [], [])
        total, per = stats[i]
        text_total.set_text(f'Remaining: {total}')
        text_floor.set_text(' '.join(f'{f}:{per[f]}' for f in range(FLOORS)))
        # 更新2D
        for f in range(FLOORS):
            fl = [(y, x) for (z, y, x) in pts if z == f]
            if fl:
                ys_f, xs_f = zip(*fl)
                scat2d[f].set_offsets(np.column_stack([xs_f, ys_f]))
            else:
                scat2d[f].set_offsets(np.empty((0, 2)))
        return [scat3d, text_total, text_floor] + scat2d

    anim = animation.FuncAnimation(fig, update, frames=len(frames), interval=300, blit=False)
    anim.save('evacuation_simulation.mp4', writer='ffmpeg', dpi=150)
    plt.close(fig)
    print('Saved evacuation_simulation.mp4')

# -----------------------------
# 9. 主程序
# -----------------------------

def main():
    env = build_environment()
    static = compute_static_field(env)
    frames, stats = simulate(env, static)
    visualize(env, frames, stats)

if __name__ == '__main__':
    main()